#!/usr/bin/python3
import rospy
from std_msgs.msg import Float64MultiArray
from interface.msg import joy_command
from hex_cfg import HexCfg
from curv_adapt_multi import CurvAdapt_multi
import torch
import random 
import os
import threading
import queue
import datetime
import time
Lock=threading.Lock()
process_deque=queue.Queue(maxsize=6)
class motor_mode:
    Disable = -1
    Zero_Torque = 0
    Pos_vel = 1
    Traj_follow = 2

def ProcessCommand(msg:joy_command):
    global device
    with Lock:
        if msg.disable_torque:
            curv_multi.set_init_done.fill_(False)
            for i in range(6):
                q_des_msg_list[i].data=[motor_mode.Disable,0,0,0,0,0,0,-999]
                q_des_pub_list[i].publish(q_des_msg_list[i])
            return

        command=torch.tensor([[msg.set_init,msg.x_vec,msg.y_vec,0,msg.w_twist]],dtype=torch.float32,device=device)
        command[:,1]=command[:,1]*hex_cfg.max_vec.x
        command[:,2]=command[:,2]*hex_cfg.max_vec.y
        command[:,3]=command[:,3]*hex_cfg.max_vec.z
        command[:,4]=command[:,4]*hex_cfg.max_vec.omega_move
        # print("command:",command)
        curv_multi.ProcessCommand(command)
        # print("curv_multi.set_init_done:",curv_multi.set_init_done)
        #use 
        # print("pub command")
        for i in range(6):
            # print("q_des[{}]={}".format(i,q_des[0,i]))
            q_pos_des=q_des[0,i,0:4].cpu().tolist()
            if curv_multi.gaits[0,i]:
                # q_pos_des[3]=-999#for stance state, set sevo disable  舵机断势能保护
                pass

            if msg.set_init==0: #2 represent traj follow mode
                q_des_msg_list[i].data=[motor_mode.Traj_follow,q_pos_des[0],q_pos_des[1],q_pos_des[2],0,0,0,q_pos_des[3]]
            else: #1 represent pos vel mode
                q_des_msg_list[i].data=[motor_mode.Pos_vel,q_pos_des[0],q_pos_des[1],q_pos_des[2],0.5,0.5,0.5,q_pos_des[3]]
                # q_des_msg_list[i].data=[motor_mode.Traj_follow,q_pos_des[0],q_pos_des[1],q_pos_des[2],0,0,0,q_pos_des[3]]
            q_des_pub_list[i].publish(q_des_msg_list[i])
    # 
            # print(leg_name_list[i],":q_des_msg=",q_des_msg_list[i].data)

        GetSuctionForce(adhesions,suction_force)
        #记录电机信息
        # if curv_multi.set_init_done.all():
        #     #t time step:
        #     q_state_dict['q_cur'].append(q_cur.clone().squeeze(0))
        #     q_state_dict['q_dot_cur'].append(q_dot_cur.clone().squeeze(0))
        #     q_state_dict['q_des'].append(q_des.clone().squeeze(0))
        #     #t-1 time step
        #     q_state_dict['torq_cur'].append(q_torq.clone().squeeze(0))
        #     if (len(q_state_dict['torq_cur'])==2000):
        #         suffix=datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        #         torch.save(q_state_dict,os.getcwd()+'/bag/motor_data/motor_raw_'+suffix+'.pt')
        #         print("save motor data to bag/motor_data/motor_raw_"+suffix+".pt")
        #         for key in q_state_dict.keys():
        #             q_state_dict[key]=[]
<<<<<<< HEAD
            # print("----------->save q_cur q_des<-------------")
            # print("q_des_t\n",q_des)
            # print("q_cur\n",q_cur)
            # print("q_dot_cur\n",q_dot_cur)
=======
        #     # print("----------->save q_cur q_des<-------------")
        #     # print("q_des_t\n",q_des)
        #     # print("q_cur\n",q_cur)
        #     # print("q_dot_cur\n",q_dot_cur)
>>>>>>> 181f11012f613eb466d33564fc88ba83963d4e05

def UpdateQState(q_state_msgs:Float64MultiArray,index):
    #实物交互更新
    global device
    # q_cur[:,index,0:3]=torch.tensor(q_state_msgs.data[0:3],dtype=torch.float32,device=device)
    # q_torq[:,index,0:3]=torch.tensor(q_state_msgs.data[3:6],dtype=torch.float32,device=device)
    # q_dot_cur[:,index,0:3]=torch.tensor(q_state_msgs.data[6:9],dtype=torch.float32,device=device)
    #仿真交互跟新
    q_cur[:,index,0:3]=torch.tensor(q_state_msgs.data[0:3],dtype=torch.float32,device=device)

def GetSuctionForce(adhesions:torch.Tensor,suction_force:torch.Tensor):
    #需要吸附
    add_force_mask=adhesions &(suction_force<=hex_cfg.max_suck_force)    
    suction_force[add_force_mask]+=100+(random.random()-0.5)*10
    suction_force[suction_force>hex_cfg.max_suck_force]=hex_cfg.max_suck_force
    #需要释放
    sub_force_mask=(~adhesions) & (suction_force>=0)
    suction_force[sub_force_mask]-=100+(random.random()-0.5)*10
    suction_force[suction_force<0]=0
    

if __name__ == '__main__':
    rospy.init_node('run_curv_adapt',anonymous=True)
    hex_cfg=HexCfg("config.yaml")
    command_sub=rospy.Subscriber('/usr/command',joy_command,ProcessCommand,queue_size=10)
    # device= 'cuda' if torch.cuda.is_available() else 'cpu'
    device='cpu'
    print("device=",device)
    q_des=torch.zeros(1,6,4,dtype=torch.float32,device=device)
    q_cur=torch.zeros(1,6,3,dtype=torch.float32,device=device)
    q_dot_cur=torch.zeros(1,6,3,dtype=torch.float32,device=device)
    q_torq=torch.zeros(1,6,3,dtype=torch.float32,device=device)
    adhesions=torch.zeros(1,6,dtype=torch.bool,device=device)
    suction_force=torch.zeros(1,6,dtype=torch.float32,device=device)
    curv_multi=CurvAdapt_multi("config.yaml",device,1,
                               q_des,q_cur,q_torq,adhesions,suction_force)
    # q_cur_flat=q_cur.view(18)
    # q_torq_flat=q_torq.view(18)
    # q_des_flat=q_des[...,0:3].view(18)
    q_des_pub_list=[]
    q_cur_sub_list=[]
    q_des_msg_list=[] #Float64MultiArray()
    q_state_dict={'q_des':[],'q_cur':[],'q_dot_cur':[],'torq_cur':[]}
    leg_name_list=['LB','LF','LM','RB','RF','RM']
    for i,leg_name in enumerate(leg_name_list):
        q_des_pub_list.append(rospy.Publisher('/'+leg_name+'/sita_des',Float64MultiArray,queue_size=10))
        q_cur_sub_list.append(rospy.Subscriber('/'+leg_name+'/sita_cur',Float64MultiArray,UpdateQState,callback_args=i))
        q_des_msg_list.append(Float64MultiArray())
    print("------------------->run curv adapt ready<-------------------")
    rospy.spin()

